{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c49a6d5",
   "metadata": {},
   "source": [
    "# Fine-tune Qwe2.5-1.5B with Alpaca Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b15eff",
   "metadata": {},
   "source": [
    "This example demonstrates how to fine-tune Qwen2.5-1.5B model with the Alpaca Dataset using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK.\n",
    "\n",
    "This notebooks walks you through the prerequisites of using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK, and how to submit TrainJob to bootstrap the fine-tuning workflow.\n",
    "\n",
    "Qwen2.5-1.5B: https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct\n",
    "\n",
    "Alpaca Dataset: https://huggingface.co/datasets/tatsu-lab/alpaca"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed4a60eb",
   "metadata": {},
   "source": [
    "## Install the Kubeflow SDK\n",
    "\n",
    "You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288ec515",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install kubeflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7211fbf9",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "### Install Official Training Runtimes\n",
    "\n",
    "You need to make sure that you've installed the Kubeflow Trainer Controller Manager and Kubeflow Training Runtimes mentioned in the [installation guide](https://www.kubeflow.org/docs/components/trainer/operator-guides/installation/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d35e4fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all available Kubeflow Training Runtimes.\n",
    "from kubeflow.trainer import *\n",
    "from kubeflow_trainer_api import models\n",
    "\n",
    "client = TrainerClient()\n",
    "for runtime in client.list_runtimes():\n",
    "    print(runtime)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb5f256",
   "metadata": {},
   "source": [
    "### Create PVCs for Models and Datasets\n",
    "\n",
    "Currently, we do not support automatically orchestrate the volume claim in (Cluster)TrainingRuntime.\n",
    "\n",
    "So, we need to manually create PVCs for each models we want to fine-tune. Please note that **the PVC name must be equal to the ClusterTrainingRuntime name**. In this example, it's `torchtune-qwen2.5-1.5b`.\n",
    "\n",
    "REF: https://github.com/kubeflow/trainer/issues/2630"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c11cc7fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a PersistentVolumeClaim for the TorchTune Qwen 2.5 1.5B model.\n",
    "client.backend.core_api.create_namespaced_persistent_volume_claim(\n",
    "    namespace=\"default\",\n",
    "    body=models.IoK8sApiCoreV1PersistentVolumeClaim(\n",
    "        apiVersion=\"v1\",\n",
    "        kind=\"PersistentVolumeClaim\",\n",
    "        metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(\n",
    "            name=\"torchtune-qwen2.5-1.5b\"\n",
    "        ),\n",
    "        spec=models.IoK8sApiCoreV1PersistentVolumeClaimSpec(\n",
    "            accessModes=[\"ReadWriteOnce\"],\n",
    "            resources=models.IoK8sApiCoreV1VolumeResourceRequirements(\n",
    "                requests={\n",
    "                    \"storage\": models.IoK8sApimachineryPkgApiResourceQuantity(\"200Gi\")\n",
    "                }\n",
    "            ),\n",
    "        ),\n",
    "    ).to_dict(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67490f66",
   "metadata": {},
   "source": [
    "## Bootstrap LLM Fine-tuning Workflow\n",
    "\n",
    "Kubeflow TrainJob will train the model in the referenced (Cluster)TrainingRuntime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641fae4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = client.train(\n",
    "    runtime=client.get_runtime(name=\"torchtune-qwen2.5-1.5b\"),\n",
    "    initializer=Initializer(\n",
    "        dataset=HuggingFaceDatasetInitializer(\n",
    "            storage_uri=\"hf://tatsu-lab/alpaca/data\"\n",
    "        ),\n",
    "        model=HuggingFaceModelInitializer(\n",
    "            storage_uri=\"hf://Qwen/Qwen2.5-1.5B-Instruct\",\n",
    "        )\n",
    "    ),\n",
    "    trainer=BuiltinTrainer(\n",
    "        config=TorchTuneConfig(\n",
    "            dataset_preprocess_config=TorchTuneInstructDataset(\n",
    "                source=DataFormat.PARQUET, split=\"train[:1000]\"\n",
    "            ),\n",
    "            resources_per_node={\n",
    "                \"memory\": \"128G\",\n",
    "                \"gpu\": 1,\n",
    "            },\n",
    "            \n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5fbe8e",
   "metadata": {},
   "source": [
    "## Wait for running status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53eaa65a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Wait for the running status.\n",
    "client.wait_for_job_status(name=job_name, status={\"Running\"})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75a82b76",
   "metadata": {},
   "source": [
    "## Watch the TrainJob Logs\n",
    "\n",
    "We can use the `get_job_logs()` API to get the TrainJob logs.\n",
    "\n",
    "### Dataset Initializer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e9d454",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kubeflow.trainer.constants import constants\n",
    "\n",
    "for line in client.get_job_logs(job_name, follow=True, step=constants.DATASET_INITIALIZER):\n",
    "    print(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f970c03",
   "metadata": {},
   "source": [
    "### Model Initializer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ce8c1df",
   "metadata": {},
   "outputs": [],
   "source": [
    "for line in client.get_job_logs(job_name, follow=True, step=constants.MODEL_INITIALIZER):\n",
    "    print(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b67775ea",
   "metadata": {},
   "source": [
    "### Trainer Node "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae672d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for c in client.get_job(name=job_name).steps:\n",
    "    print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")\n",
    "\n",
    "for line in client.get_job_logs(job_name, follow=True):\n",
    "    print(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "042ebbb6",
   "metadata": {},
   "source": [
    "# Get the Fine-tuned Model\n",
    "\n",
    "After Trainer node completes the fine-tuning task, the fine-tuned model will be stored into the `/workspace/output` directory, which can be shared across Pods through PVC mounting. You can find it in another Pod's `/<mountDir>/output` directory if you mount the PVC under `/<mountDir>`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
